{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 虚拟环境必须 python 3.9+(因为whisper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cell_style": "center",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9EqJBfYhSvIt",
    "outputId": "a23f1487-20fb-4f3c-873f-d620b29c448a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/104.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m104.4/104.4 kB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Building wheel for pysrt (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
     ]
    }
   ],
   "source": [
    "!pip install -q pysrt\n",
    "!pip install -q pysubs2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KSY40g3JSZDp"
   },
   "source": [
    "### 1.根据字幕切割出分段音频"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "43RQDWgNSlFz"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import subprocess\n",
    "from collections import Counter\n",
    "\n",
    "import chardet\n",
    "import pysrt\n",
    "import pysubs2\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zoXRJDUkTQJA"
   },
   "source": [
    "#### 常用函数\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "huizrd80TENm"
   },
   "outputs": [],
   "source": [
    "def get_subdir(directory):\n",
    "    subdirectories = []\n",
    "    for dirpath, dirnames, files in os.walk(directory):\n",
    "        for dirname in dirnames:\n",
    "            subdirectories.append(os.path.join(dirpath, dirname))\n",
    "    subdirectories.sort()\n",
    "    return subdirectories\n",
    "\n",
    "def get_filename(directory,format=None):\n",
    "    file_list = []\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if format:\n",
    "                if file.endswith(format):\n",
    "                    file_path = os.path.join(root, file)\n",
    "                    file_list.append([file,file_path])\n",
    "            else:\n",
    "                file_path = os.path.join(root, file)\n",
    "                file_list.append([file, file_path])\n",
    "    file_list.sort()\n",
    "    return file_list\n",
    "\n",
    "\n",
    "#获取一级子目录\n",
    "def get_first_subdir(directory):\n",
    "    subdirectories = []\n",
    "    for name in os.listdir(directory):\n",
    "        if os.path.isdir(os.path.join(directory, name)):\n",
    "            subdirectories.append(os.path.join(directory, name))\n",
    "    subdirectories.sort()\n",
    "    return subdirectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "XULhgXcLTH6H"
   },
   "outputs": [],
   "source": [
    "def detect_encoding(file_name):\n",
    "    with open(file_name, 'rb') as file:\n",
    "        result = chardet.detect(file.read())\n",
    "    return result['encoding']\n",
    "\n",
    "def most_common_element(lst,num=1):\n",
    "    counter = Counter(lst)\n",
    "    most = counter.most_common(num)\n",
    "    return most\n",
    "\n",
    "\n",
    "def make_filename_safe(filename):\n",
    "    # 将非法字符替换为下划线\n",
    "    filename = re.sub(r'[\\\\/:*?\"<>|]', '_', filename)\n",
    "    # 去除多余的空格\n",
    "    filename = re.sub(r'\\s+', ' ', filename)\n",
    "    # 去除开头和结尾的空格\n",
    "    filename = filename.strip()\n",
    "    return filename\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vimRKLB1TXXp"
   },
   "source": [
    "#### VideoSegmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "OePQQ39ARtdv"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class VideoSegmentation:\n",
    "    def __init__(self, video_lis_pth,audio_out_dir,subtitle_dir):\n",
    "        self.video_lis_pth = video_lis_pth\n",
    "        self.audio_out_dir = audio_out_dir\n",
    "        self.subtitle_dir = subtitle_dir\n",
    "\n",
    "\n",
    "    def process(self):\n",
    "        video_lis = get_filename(self.video_lis_pth)\n",
    "\n",
    "        style = ''\n",
    "        sub_format = ''\n",
    "        voice_dir = 'voice'\n",
    "        for file, pth in tqdm(video_lis[:], desc='Processing Videos'):\n",
    "\n",
    "            filename, format = os.path.splitext(file)\n",
    "            # 创建对应的音频文件夹\n",
    "            os.makedirs(f'{self.audio_out_dir}/{filename}', exist_ok=True)\n",
    "            os.makedirs(f'{self.audio_out_dir}/{filename}/{voice_dir}', exist_ok=True)\n",
    "\n",
    "            if self.subtitle_dir:\n",
    "                if not sub_format:\n",
    "                    # 选择一个字幕文件 获取字幕文件的格式和编码\n",
    "                    one_subtitle_file = os.path.join(self.subtitle_dir,os.listdir(self.subtitle_dir)[0])\n",
    "                    sub_file,sub_format = os.path.splitext(one_subtitle_file)\n",
    "                    encoding = detect_encoding(one_subtitle_file)\n",
    "\n",
    "                # 获取当前视频对应的字幕文件\n",
    "                cur_sub_file = f'{self.subtitle_dir}/{filename}{sub_format}'\n",
    "                # 获取对应字幕\n",
    "                if sub_format == '.srt':\n",
    "\n",
    "                    srt_file = pysrt.open(cur_sub_file, encoding=encoding)\n",
    "                    for index, subtitle in enumerate(srt_file[:]):\n",
    "                        # 获取开始和结束时间\n",
    "\n",
    "                        start_time = subtitle.start\n",
    "                        end_time = subtitle.end\n",
    "\n",
    "                        start_time = start_time.to_time()\n",
    "                        end_time = end_time.to_time()\n",
    "                        # print(f'开始时间：{start_time}，结束时间：{end_time}')\n",
    "\n",
    "                        # 使用FFmpeg切割视频 改成mp3就无法输出\n",
    "                        audio_output = f'{self.audio_out_dir}/{filename}/{voice_dir}/{index}_{make_filename_safe(subtitle.text)}.wav'\n",
    "\n",
    "                        command = ['ffmpeg', '-ss', str(start_time), '-to', str(end_time), '-i', f'{pth}', \"-vn\",  '-c:a', 'pcm_s16le',\n",
    "                                         audio_output,  '-loglevel', 'quiet']\n",
    "\n",
    "                        subprocess.run(command)\n",
    "                elif sub_format == '.ass':\n",
    "                    subs = pysubs2.load(cur_sub_file, encoding=encoding)\n",
    "                    if not style:\n",
    "                        style_lis = [sub.style for sub in subs]\n",
    "                        most_1 = most_common_element(style_lis)\n",
    "                        style = most_1[0][0]\n",
    "\n",
    "                    new_subs = [sub for sub in subs if sub.style == style]\n",
    "                    for index, subtitle in enumerate(new_subs[:]):\n",
    "                        # 获取开始和结束时间\n",
    "                        if subtitle.style == style:\n",
    "                            start_time = subtitle.start\n",
    "                            end_time = subtitle.end\n",
    "\n",
    "                            start_time = start_time / 1000\n",
    "                            end_time = end_time / 1000\n",
    "\n",
    "\n",
    "                            # 使用FFmpeg切割视频 改成mp3就无法输出\n",
    "                            # audio_output = f'{self.audio_out_dir}/{filename}/{index}_{make_filename_safe(subtitle.text)}_ass.wav'\n",
    "                            audio_output = f'{self.audio_out_dir}/{filename}/{voice_dir}/{index}_{make_filename_safe(subtitle.text)}.wav'\n",
    "\n",
    "                            command = ['ffmpeg', '-ss', str(start_time), '-to', str(end_time), '-i', f'{pth}', \"-vn\",  '-c:a', 'pcm_s16le',\n",
    "                                             audio_output,  '-loglevel', 'quiet']\n",
    "\n",
    "                            subprocess.run(command)\n",
    "\n",
    "        exit()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zN29Mm7oT54t"
   },
   "source": [
    "#### 自定义config参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "n77bOc_hT9md"
   },
   "outputs": [],
   "source": [
    "# 自定义修改下面字典的value值  标有**的文件夹，需要有对应文件。\n",
    "video_config = {\"video_lis_pth\": \"/content/drive/MyDrive/GPTData/origin_video\",  # **视频所在文件夹 **需要准备视\n",
    "                }\n",
    "audio_config = {\n",
    "                \"audio_model_pth\":'/content/drive/MyDrive/GPTData/voicemodel',  # **模型权重路径 需要下载模型→git clone https://huggingface.co/scixing/voicemodel\n",
    "                \"audio_roles_dir\":'/content/drive/MyDrive/GPTData/roles', # **分类好的角色音频路径 需要手动分类\n",
    "                \"audio_out_dir\": \"/content/drive/MyDrive/GPTData/audio\",  # 视频切割输出的音频路径\n",
    "                }\n",
    "\n",
    "srt_config = {\n",
    "               \"subtitle_dir\":\"/content/drive/MyDrive/GPTData/srt\",  # **视频对应字幕，视频和字幕名称需要一致 需要准备 ,\n",
    "                \"srt_out_dir\":\"/content/drive/MyDrive/GPTData/roletxt\",  # 预测的角色类型路径\n",
    "            }\n",
    "\n",
    "\n",
    "# whisper 模型和下载链接，不建议使用 tiny 和large 模型\n",
    "WHISPER_MODELS_LIST = ['tiny.en', 'tiny', 'base.en', 'base', 'small.en', 'small', \n",
    "                            'medium.en', 'medium', 'large-v1', 'large-v2', 'large']\n",
    "WHISPER_MODELS_LINKS = {\n",
    "    \"tiny.en\": \"https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt\",\n",
    "    \"tiny\": \"https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt\",\n",
    "    \"base.en\": \"https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt\",\n",
    "    \"base\": \"https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt\",\n",
    "    \"small.en\": \"https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt\",\n",
    "    \"small\": \"https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt\",\n",
    "    \"medium.en\": \"https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt\",\n",
    "    \"medium\": \"https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt\",\n",
    "    \"large-v1\": \"https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt\",\n",
    "    \"large-v2\": \"https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt\",\n",
    "    \"large\": \"https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  Video to Subtitles with Whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install openai-whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt\n"
     ]
    }
   ],
   "source": [
    "print(WHISPER_MODELS_LINKS[\"small\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget 模型地址"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ffmpeg\n",
    "import tempfile\n",
    "import torch\n",
    "import whisper\n",
    "from whisper.utils import get_writer\n",
    "from typing import Iterator, TextIO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Video2Subtitles(object):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def srt_format_timestamp(self, seconds: float):\n",
    "        assert seconds >= 0, \"non-negative timestamp expected\"\n",
    "        milliseconds = round(seconds * 1000.0)\n",
    "\n",
    "        hours = milliseconds // 3_600_000\n",
    "        milliseconds -= hours * 3_600_000\n",
    "\n",
    "        minutes = milliseconds // 60_000\n",
    "        milliseconds -= minutes * 60_000\n",
    "\n",
    "        seconds = milliseconds // 1_000\n",
    "        milliseconds -= seconds * 1_000\n",
    "\n",
    "        return (f\"{hours}:\") + f\"{minutes:02d}:{seconds:02d},{milliseconds:03d}\"\n",
    "\n",
    "    def write_srt(self, transcript: Iterator[dict], file: TextIO):\n",
    "        count = 0\n",
    "        for segment in transcript:\n",
    "            count += 1\n",
    "            print(\n",
    "                # f\"{count}\\n\"\n",
    "                f\"{self.srt_format_timestamp(segment['start'])} --> {self.srt_format_timestamp(segment['end'])}\\n\"\n",
    "                f\"{segment['text'].replace('-->', '->').strip()}\\n\",\n",
    "                file=file,\n",
    "                flush=True,\n",
    "            )\n",
    "\n",
    "    def transcribe(self, input_video: str, lang: str, MODEL_WHISPER: str, task: str, subtitle_format: str,\n",
    "                   AddSrtToVideo: bool):\n",
    "        \"\"\"\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_video:\n",
    "        lang: language of your input file\n",
    "        MODEL_WHISPER: tiny/small /base/large,you can also download into your local path  eg. /tiny.pt\n",
    "        task: transcribe/translate(any language to english)\n",
    "        subtitle_format:\"txt\", \"vtt\", \"srt\", \"tsv\",  \"json\",\n",
    "        AddSrtToVideo:\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        }\n",
    "        \"\"\"\n",
    "        DEVICE = torch.cuda.is_available()\n",
    "        model = whisper.load_model(MODEL_WHISPER)\n",
    "        input_video_ = input_video if isinstance(input_video, str) else input_video.name\n",
    "        result = model.transcribe(\n",
    "            input_video_,\n",
    "            task=task,\n",
    "            language=lang,\n",
    "            verbose=True,\n",
    "            initial_prompt=None,\n",
    "            word_timestamps=False,\n",
    "            fp16=DEVICE\n",
    "        )\n",
    "        subtitle_file = input_video_.rsplit(\".\", 1)[0] + \".\" + subtitle_format\n",
    "        print(\"subtitle_file:\", subtitle_file)\n",
    "        writer = get_writer(subtitle_format, str(tempfile.gettempdir()))\n",
    "        writer(result, subtitle_file)\n",
    "        if subtitle_format == \"srt\":\n",
    "            with open(subtitle_file, \"w\") as srt:\n",
    "                self.write_srt(result[\"segments\"], file=srt)\n",
    "        if AddSrtToVideo:\n",
    "            return self.add_srt_to_video(input_video_, subtitle_file)\n",
    "        return subtitle_file\n",
    "\n",
    "    def add_srt_to_video(self, input_video_, subtitle_file):\n",
    "        video_out = input_video_ + \"_output.mp4\"\n",
    "        input_ffmpeg = ffmpeg.input(input_video_)\n",
    "        input_ffmpeg_sub = ffmpeg.input(subtitle_file)\n",
    "        input_video = input_ffmpeg['v']\n",
    "        input_audio = input_ffmpeg['a']\n",
    "        input_subtitles = input_ffmpeg_sub['s']\n",
    "        stream = ffmpeg.output(\n",
    "            input_video, input_audio, input_subtitles, video_out,\n",
    "            vcodec='copy', acodec='copy', scodec=subtitle_format\n",
    "        )\n",
    "        stream = ffmpeg.overwrite_output(stream)\n",
    "        ffmpeg.run(stream)\n",
    "        return video_out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 运行: Video to Subtitles with Whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_video = video_config[\"video_lis_pth\"]+\"/input.mp4\"\n",
    "MODEL_WHISPER = \"small\" # or your downloaded model path\"\n",
    "lang = \"zh\" # 输入视频的语言， 可支持各种语言 \n",
    "task = \"transcribe\" # translate or transcribe\n",
    "AddSrtToVideo = False\n",
    "subtitle_format = \"srt\"\n",
    "Video2Subtitles().transcribe(input_video, lang, MODEL_WHISPER, task, subtitle_format, AddSrtToVideo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9PgLSO5OTlib"
   },
   "source": [
    "#### 运行1 音频提取分割"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "_LQVJAZTSWsh"
   },
   "outputs": [],
   "source": [
    "# video_segmentor = VideoSegmentation(video_config['video_lis_pth'],\n",
    "#                                         audio_config['audio_out_dir'],\n",
    "#                                         srt_config['subtitle_dir'])\n",
    "# video_segmentor.process()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "oIWUNrcukmGw"
   },
   "outputs": [],
   "source": [
    "# !git clone https://huggingface.co/scixing/voicemodel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "45LvR5uOnt1w"
   },
   "source": [
    "### 2.音频特征提取\n",
    "Audio Feature Extraction  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I_I36FzElbVv"
   },
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "JT45f9PCiA-g"
   },
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import pickle\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Parameter\n",
    "import random\n",
    "import sys\n",
    "from datetime import datetime\n",
    "import librosa\n",
    "from torch.utils import data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TTwWayPi2qMV"
   },
   "source": [
    "#### Audio模型定义辅助函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "uTiSNi9UwIK2"
   },
   "outputs": [],
   "source": [
    "def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):\n",
    "    return nn.Sequential(\n",
    "        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),\n",
    "        Res2Conv1dReluBn(channels, kernel_size, stride, padding, dilation, scale=scale),\n",
    "        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),\n",
    "        SE_Connect(channels)\n",
    "    )\n",
    "\n",
    "\n",
    "def load_audio(audio_path,\n",
    "               feature_method='melspectrogram',\n",
    "               mode='train',\n",
    "               sr=16000,\n",
    "               chunk_duration=3,\n",
    "               min_duration=0.5,\n",
    "               augmentors=None):\n",
    "    \"\"\"\n",
    "    加载并预处理音频\n",
    "    :param audio_path: 音频路径\n",
    "    :param feature_method: 预处理方法\n",
    "    :param mode: 对数据处理的方式，包括train，eval，infer\n",
    "    :param sr: 采样率\n",
    "    :param chunk_duration: 训练或者评估使用的音频长度\n",
    "    :param min_duration: 最小训练或者评估的音频长度\n",
    "    :param augmentors: 数据增强方法\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # 读取音频数据\n",
    "    wav, sr_ret = librosa.load(audio_path, sr=sr)\n",
    "    num_wav_samples = wav.shape[0]\n",
    "    # 数据太短不利于训练\n",
    "    if mode == 'train':\n",
    "        if num_wav_samples < int(min_duration * sr):\n",
    "            raise Exception(f'音频长度小于{min_duration}s，实际长度为：{(num_wav_samples / sr):.2f}s')\n",
    "    # 对小于训练长度的复制补充\n",
    "    num_chunk_samples = int(chunk_duration * sr)\n",
    "    if num_wav_samples <= num_chunk_samples:\n",
    "        shortage = num_chunk_samples - num_wav_samples\n",
    "        wav = np.pad(wav, (0, shortage), 'wrap')\n",
    "    # 裁剪需要的数据\n",
    "    if mode == 'train':\n",
    "        # 随机裁剪\n",
    "        num_wav_samples = wav.shape[0]\n",
    "        num_chunk_samples = int(chunk_duration * sr)\n",
    "        if num_wav_samples > num_chunk_samples + 1:\n",
    "            start = random.randint(0, num_wav_samples - num_chunk_samples - 1)\n",
    "            stop = start + num_chunk_samples\n",
    "            wav = wav[start:stop]\n",
    "            # 对每次都满长度的再次裁剪\n",
    "            if random.random() > 0.5:\n",
    "                wav[:random.randint(1, sr // 4)] = 0\n",
    "                wav = wav[:-random.randint(1, sr // 4)]\n",
    "        # 数据增强\n",
    "        if augmentors is not None:\n",
    "            for key, augmentor in augmentors.items():\n",
    "                if key == 'specaug': continue\n",
    "                wav = augmentor(wav)\n",
    "    elif mode == 'eval':\n",
    "        # 为避免显存溢出，只裁剪指定长度\n",
    "        num_wav_samples = wav.shape[0]\n",
    "        num_chunk_samples = int(chunk_duration * sr)\n",
    "        if num_wav_samples > num_chunk_samples + 1:\n",
    "            wav = wav[:num_chunk_samples]\n",
    "    # 获取音频特征\n",
    "    if feature_method == 'melspectrogram':\n",
    "        # 计算梅尔频谱\n",
    "        features = librosa.feature.melspectrogram(y=wav, sr=sr, n_fft=400, n_mels=80, hop_length=160, win_length=400)\n",
    "    elif feature_method == 'spectrogram':\n",
    "        # 计算声谱图\n",
    "        linear = librosa.stft(wav, n_fft=400, win_length=400, hop_length=160)\n",
    "        features, _ = librosa.magphase(linear)\n",
    "    else:\n",
    "        raise Exception(f'预处理方法 {feature_method} 不存在！')\n",
    "    features = librosa.power_to_db(features, ref=1.0, amin=1e-10, top_db=None)\n",
    "    # 数据增强\n",
    "    if mode == 'train' and augmentors is not None:\n",
    "        for key, augmentor in augmentors.items():\n",
    "            if key == 'specaug':\n",
    "                features = augmentor(features)\n",
    "    # 归一化\n",
    "    mean = np.mean(features, 0, keepdims=True)\n",
    "    std = np.std(features, 0, keepdims=True)\n",
    "    features = (features - mean) / (std + 1e-5)\n",
    "    return features\n",
    "\n",
    "\n",
    "class Res2Conv1dReluBn(nn.Module):\n",
    "    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False, scale=4):\n",
    "        super().__init__()\n",
    "        assert channels % scale == 0, \"{} % {} != 0\".format(channels, scale)\n",
    "        self.scale = scale\n",
    "        self.width = channels // scale\n",
    "        self.nums = scale if scale == 1 else scale - 1\n",
    "\n",
    "        self.convs = []\n",
    "        self.bns = []\n",
    "        for i in range(self.nums):\n",
    "            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))\n",
    "            self.bns.append(nn.BatchNorm1d(self.width))\n",
    "        self.convs = nn.ModuleList(self.convs)\n",
    "        self.bns = nn.ModuleList(self.bns)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = []\n",
    "        spx = torch.split(x, self.width, 1)\n",
    "        for i in range(self.nums):\n",
    "            if i == 0:\n",
    "                sp = spx[i]\n",
    "            else:\n",
    "                sp = sp + spx[i]\n",
    "            # Order: conv -> relu -> bn\n",
    "            sp = self.convs[i](sp)\n",
    "            sp = self.bns[i](F.relu(sp))\n",
    "            out.append(sp)\n",
    "        if self.scale != 1:\n",
    "            out.append(spx[self.nums])\n",
    "        out = torch.cat(out, dim=1)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Conv1dReluBn(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)\n",
    "        self.bn = nn.BatchNorm1d(out_channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.bn(F.relu(self.conv(x)))\n",
    "\n",
    "\n",
    "class SE_Connect(nn.Module):\n",
    "    def __init__(self, channels, s=2):\n",
    "        super().__init__()\n",
    "        assert channels % s == 0, \"{} % {} != 0\".format(channels, s)\n",
    "        self.linear1 = nn.Linear(channels, channels // s)\n",
    "        self.linear2 = nn.Linear(channels // s, channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = x.mean(dim=2)\n",
    "        out = F.relu(self.linear1(out))\n",
    "        out = torch.sigmoid(self.linear2(out))\n",
    "        out = x * out.unsqueeze(2)\n",
    "        return out\n",
    "\n",
    "\n",
    "class AttentiveStatsPool(nn.Module):\n",
    "    def __init__(self, in_dim, bottleneck_dim):\n",
    "        super().__init__()\n",
    "        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.\n",
    "        self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1)  # equals W and b in the paper\n",
    "        self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1)  # equals V and k in the paper\n",
    "\n",
    "    def forward(self, x):\n",
    "        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.\n",
    "        alpha = torch.tanh(self.linear1(x))\n",
    "        alpha = torch.softmax(self.linear2(alpha), dim=2)\n",
    "        mean = torch.sum(alpha * x, dim=2)\n",
    "        residuals = torch.sum(alpha * x ** 2, dim=2) - mean ** 2\n",
    "        std = torch.sqrt(residuals.clamp(min=1e-9))\n",
    "        return torch.cat([mean, std], dim=1)\n",
    "\n",
    "\n",
    "class EcapaTdnn(nn.Module):\n",
    "    def __init__(self, input_size=80, channels=512, embd_dim=192):\n",
    "        super().__init__()\n",
    "        self.layer1 = Conv1dReluBn(input_size, channels, kernel_size=5, padding=2, dilation=1)\n",
    "        self.layer2 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=2, dilation=2, scale=8)\n",
    "        self.layer3 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=3, dilation=3, scale=8)\n",
    "        self.layer4 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=4, dilation=4, scale=8)\n",
    "\n",
    "        cat_channels = channels * 3\n",
    "        out_channels = cat_channels * 2\n",
    "        self.emb_size = embd_dim\n",
    "        self.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1)\n",
    "        self.pooling = AttentiveStatsPool(cat_channels, 128)\n",
    "        self.bn1 = nn.BatchNorm1d(out_channels)\n",
    "        self.linear = nn.Linear(out_channels, embd_dim)\n",
    "        self.bn2 = nn.BatchNorm1d(embd_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out1 = self.layer1(x)\n",
    "        out2 = self.layer2(out1) + out1\n",
    "        out3 = self.layer3(out1 + out2) + out1 + out2\n",
    "        out4 = self.layer4(out1 + out2 + out3) + out1 + out2 + out3\n",
    "\n",
    "        out = torch.cat([out2, out3, out4], dim=1)\n",
    "        out = F.relu(self.conv(out))\n",
    "        out = self.bn1(self.pooling(out))\n",
    "        out = self.bn2(self.linear(out))\n",
    "        return out\n",
    "\n",
    "\n",
    "class SpeakerIdetification(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            backbone,\n",
    "            num_class=1,\n",
    "            lin_blocks=0,\n",
    "            lin_neurons=192,\n",
    "            dropout=0.1, ):\n",
    "        \"\"\"The speaker identification model, which includes the speaker backbone network\n",
    "           and the a linear transform to speaker class num in training\n",
    "\n",
    "        Args:\n",
    "            backbone (Paddle.nn.Layer class): the speaker identification backbone network model\n",
    "            num_class (_type_): the speaker class num in the training dataset\n",
    "            lin_blocks (int, optional): the linear layer transform between the embedding and the final linear layer. Defaults to 0.\n",
    "            lin_neurons (int, optional): the output dimension of final linear layer. Defaults to 192.\n",
    "            dropout (float, optional): the dropout factor on the embedding. Defaults to 0.1.\n",
    "        \"\"\"\n",
    "        super(SpeakerIdetification, self).__init__()\n",
    "        # speaker idenfication backbone network model\n",
    "        # the output of the backbond network is the target embedding\n",
    "        self.backbone = backbone\n",
    "        if dropout > 0:\n",
    "            self.dropout = nn.Dropout(dropout)\n",
    "        else:\n",
    "            self.dropout = None\n",
    "\n",
    "        # construct the speaker classifer\n",
    "        input_size = self.backbone.emb_size\n",
    "        self.blocks = list()\n",
    "        for i in range(lin_blocks):\n",
    "            self.blocks.extend([\n",
    "                nn.BatchNorm1d(input_size),\n",
    "                nn.Linear(in_features=input_size, out_features=lin_neurons),\n",
    "            ])\n",
    "            input_size = lin_neurons\n",
    "\n",
    "        # the final layer\n",
    "        self.weight = Parameter(torch.FloatTensor(num_class, input_size), requires_grad=True)\n",
    "        nn.init.xavier_normal_(self.weight, gain=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Do the speaker identification model forwrd,\n",
    "           including the speaker embedding model and the classifier model network\n",
    "\n",
    "        Args:\n",
    "            x (paddle.Tensor): input audio feats,\n",
    "                               shape=[batch, dimension, times]\n",
    "            lengths (paddle.Tensor, optional): input audio length.\n",
    "                                        shape=[batch, times]\n",
    "                                        Defaults to None.\n",
    "\n",
    "        Returns:\n",
    "            paddle.Tensor: return the logits of the feats\n",
    "        \"\"\"\n",
    "        # x.shape: (N, C, L)\n",
    "        x = self.backbone(x)  # (N, emb_size)\n",
    "        if self.dropout is not None:\n",
    "            x = self.dropout(x)\n",
    "\n",
    "        for fc in self.blocks:\n",
    "            x = fc(x)\n",
    "\n",
    "        logits = F.linear(F.normalize(x), F.normalize(self.weight, dim=-1))\n",
    "\n",
    "        return logits\n",
    "\n",
    "\n",
    "# 数据加载器\n",
    "class CustomDataset(data.Dataset):\n",
    "    \"\"\"\n",
    "    加载并预处理音频\n",
    "    :param data_list_path: 数据列表\n",
    "    :param feature_method: 预处理方法\n",
    "    :param mode: 对数据处理的方式，包括train，eval，infer\n",
    "    :param sr: 采样率\n",
    "    :param chunk_duration: 训练或者评估使用的音频长度\n",
    "    :param min_duration: 最小训练或者评估的音频长度\n",
    "    :param augmentors: 数据增强方法\n",
    "    :return:\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_list_path,\n",
    "                 feature_method='melspectrogram',\n",
    "                 mode='train',\n",
    "                 sr=16000,\n",
    "                 chunk_duration=3,\n",
    "                 min_duration=0.5,\n",
    "                 augmentors=None):\n",
    "        super(CustomDataset, self).__init__()\n",
    "        # 当预测时不需要获取数据\n",
    "        if data_list_path is not None:\n",
    "            with open(data_list_path, 'r') as f:\n",
    "                self.lines = f.readlines()\n",
    "        self.feature_method = feature_method\n",
    "        self.mode = mode\n",
    "        self.sr = sr\n",
    "        self.chunk_duration = chunk_duration\n",
    "        self.min_duration = min_duration\n",
    "        self.augmentors = augmentors\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        try:\n",
    "            audio_path, label = self.lines[idx].replace('\\n', '').split('\\t')\n",
    "            # 加载并预处理音频\n",
    "            features = load_audio(audio_path, feature_method=self.feature_method, mode=self.mode, sr=self.sr,\n",
    "                                  chunk_duration=self.chunk_duration, min_duration=self.min_duration,\n",
    "                                  augmentors=self.augmentors)\n",
    "            return features, np.array(int(label), dtype=np.int64)\n",
    "        except Exception as ex:\n",
    "            print(f\"[{datetime.now()}] 数据: {self.lines[idx]} 出错，错误信息: {ex}\", file=sys.stderr)\n",
    "            rnd_idx = np.random.randint(self.__len__())\n",
    "            return self.__getitem__(rnd_idx)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.lines)\n",
    "\n",
    "    @property\n",
    "    def input_size(self):\n",
    "        if self.feature_method == 'melspectrogram':\n",
    "            return 80\n",
    "        elif self.feature_method == 'spectrogram':\n",
    "            return 201\n",
    "        else:\n",
    "            raise Exception(f'预处理方法 {self.feature_method} 不存在！')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JLiuN00s2lJr"
   },
   "source": [
    "#### AudioFeatureExtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "pgwnzX6p1Z8m"
   },
   "outputs": [],
   "source": [
    "class AudioFeatureExtraction:\n",
    "    def __init__(self, model_local_pth, audio_duration=3, feature_method='melspectrogram', ):\n",
    "        self.use_model = ''\n",
    "        self.audio_duration = audio_duration\n",
    "        self.feature_method = feature_method\n",
    "        self.resume = model_local_pth\n",
    "        self.model = None\n",
    "        self.device = None\n",
    "        self.load_model()\n",
    "\n",
    "    def load_model(self):\n",
    "        dataset = CustomDataset(data_list_path=None, feature_method=self.feature_method)\n",
    "        ecapa_tdnn = EcapaTdnn(input_size=dataset.input_size)\n",
    "        self.model = SpeakerIdetification(backbone=ecapa_tdnn)\n",
    "        self.device = torch.device(\"cuda\")\n",
    "        self.model.to(self.device)\n",
    "\n",
    "        # 加载模型\n",
    "        model_path = os.path.join(self.resume, self.use_model, 'model.pth')\n",
    "        model_dict = self.model.state_dict()\n",
    "        param_state_dict = torch.load(model_path)\n",
    "        for name, weight in model_dict.items():\n",
    "            if name in param_state_dict.keys():\n",
    "                if list(weight.shape) != list(param_state_dict[name].shape):\n",
    "                    param_state_dict.pop(name, None)\n",
    "        self.model.load_state_dict(param_state_dict, strict=False)\n",
    "        print(f\"成功加载模型参数和优化方法参数：{model_path}\")\n",
    "        self.model.eval()\n",
    "\n",
    "    def infer(self, audio_path):\n",
    "        data = load_audio(audio_path, mode='infer', feature_method=self.feature_method,\n",
    "                          chunk_duration=self.audio_duration)\n",
    "        data = data[np.newaxis, :]\n",
    "        data = torch.tensor(data, dtype=torch.float32, device=self.device)\n",
    "        feature = self.model.backbone(data)\n",
    "        return feature.data.cpu().numpy()\n",
    "\n",
    "    def extract_features(self, root_dir):\n",
    "        sub_dirs = get_subdir(root_dir)\n",
    "\n",
    "        for dir in sub_dirs:\n",
    "            voice_files = get_filename(os.path.join(dir, 'voice'))\n",
    "            for file, pth in voice_files:\n",
    "                new_dir = os.path.join(dir, 'feature')\n",
    "                os.makedirs(new_dir, exist_ok=True)\n",
    "                feature = self.infer(pth)[0]\n",
    "                with open(f\"{new_dir}/{file}.pkl\", \"wb\") as f:\n",
    "                    pickle.dump(feature, f)\n",
    "        print('音频特征提取完成')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sr4PIZVA2yr7"
   },
   "source": [
    "#### 运行2 音频embedding生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dEtTsba823EL",
    "outputId": "b680d1a5-d5c4-4dde-eca1-4fdf6f2c539a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "成功加载模型参数和优化方法参数：/content/drive/MyDrive/GPTData/voicemodel/model.pth\n",
      "音频特征提取完成\n"
     ]
    }
   ],
   "source": [
    "# 模型参数在第一部分 自定义config参数\n",
    "# audio_feature_extractor = AudioFeatureExtraction(audio_config['audio_model_pth'])\n",
    "# audio_feature_extractor.extract_features(audio_config['audio_out_dir'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s-JM9J4N6nDK"
   },
   "source": [
    "### 识别台本角色"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nHgJU3Sp-Bvm"
   },
   "source": [
    "#### 导包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "AWLkvOdT6OPK"
   },
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.model_selection import cross_val_score\n",
    "from scipy.spatial.distance import cosine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lA3hf-vj96-i"
   },
   "source": [
    "#### My_Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "S749oqiV9q7v"
   },
   "outputs": [],
   "source": [
    "class My_Classifier:\n",
    "    def __init__(self, feature, labels):\n",
    "        self.feature = feature\n",
    "        self.labels = labels\n",
    "\n",
    "    def predict(self, x):\n",
    "        min_dist = float('inf')\n",
    "        predicted_label = None\n",
    "\n",
    "        for i, f in enumerate(self.feature):\n",
    "            dist = cosine(x, f)\n",
    "            if dist < min_dist:\n",
    "                min_dist = dist\n",
    "                predicted_label = self.labels[i]\n",
    "\n",
    "        return predicted_label, min_dist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v1KykX3V-AAP"
   },
   "source": [
    "#### AudioClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "IdQDjE5s9oFN"
   },
   "outputs": [],
   "source": [
    "class AudioClassification:\n",
    "    def __init__(self, audio_roles_dir,srt_out_dir,audio_out_dir):\n",
    "        self.audio_roles_dir = audio_roles_dir\n",
    "        self.srt_out_dir = srt_out_dir\n",
    "\n",
    "        self.audio_first_dir = get_first_subdir(audio_out_dir)\n",
    "        self.candidate_path = self.audio_first_dir[:]\n",
    "\n",
    "        self.roles, self.roles_list = self.get_roles_list()\n",
    "        self.features, self.labels = self.get_features()\n",
    "        self.feat_sel, self.label_sel = self.get_feat_sel()\n",
    "        self.my_classifier = My_Classifier(self.feat_sel, self.label_sel)\n",
    "\n",
    "\n",
    "\n",
    "    def get_roles_list(self):\n",
    "        roles = os.listdir(self.audio_roles_dir)\n",
    "\n",
    "        \"\"\"\n",
    "        roles\n",
    "            春日 谷口 长门 新川先生 多丸裕 朝比奈 朝仓 ...\n",
    "        roles_list\n",
    "            ['209_或者这座岛有没有被当地人称为「什么什么岛」的传闻？.wav',\n",
    "            '106_我要喝100%果汁.wav', '121_你睡什么觉啊  笨蛋.wav',\n",
    "        \"\"\"\n",
    "        roles_list = []\n",
    "        roles_list_full = []\n",
    "        for role in roles:\n",
    "\n",
    "            sub_list = os.listdir(os.path.join(self.audio_roles_dir,role))\n",
    "            roles_list.append(sub_list)\n",
    "\n",
    "            full_name_list = [self.audio_roles_dir+role+'/'+file for file in sub_list]\n",
    "            roles_list_full.append(full_name_list)\n",
    "\n",
    "        return roles, roles_list\n",
    "\n",
    "\n",
    "\n",
    "    def get_features(self):\n",
    "        features = []\n",
    "        labels = []\n",
    "        dim = 0\n",
    "        count = 0\n",
    "\n",
    "        for role in self.roles:\n",
    "            print(role,end='')\n",
    "            for file in self.roles_list[self.roles.index(role)]:\n",
    "                deal_flag = False\n",
    "                for candidate in self.candidate_path: #'/mnt/sda/baidu_disk/lg/scixing/Haruhi ep1'\n",
    "                    candidate_fname = os.path.join(candidate, 'voice')\n",
    "                    if os.path.exists(candidate_fname):\n",
    "                        deal_flag = True\n",
    "                        feature_fname = os.path.join(candidate,'feature',file) +'.pkl'\n",
    "                        break\n",
    "\n",
    "                if deal_flag == False:\n",
    "                    # print('warning!', file, 'not found')\n",
    "                    continue\n",
    "\n",
    "                if not os.path.exists(feature_fname):\n",
    "                    # print('warning!', feature_fname, 'not found')\n",
    "                    continue\n",
    "\n",
    "                    # pinkle load feature_fname\n",
    "                with open(feature_fname, 'rb') as f:\n",
    "                    feature = pickle.load(f)\n",
    "\n",
    "                count += 1\n",
    "\n",
    "                # append numpy array feature into numpy matrix features\n",
    "                if dim == 0:\n",
    "                    features = feature\n",
    "                    dim = feature.shape[0]\n",
    "                    # print(dim)\n",
    "                else:\n",
    "                    features = np.vstack((features, feature))\n",
    "\n",
    "                labels.append(role)\n",
    "\n",
    "                # print(feature_fname,'found')\n",
    "\n",
    "            # break\n",
    "        return features, labels\n",
    "\n",
    "    def knn_test(self):\n",
    "        \"\"\"\n",
    "        feature是一个N*D的numpy矩阵，每行存储了一个D维特征 labels是一个python的list of string，表示每行对应的数据的标签。\n",
    "        我想验证这批数据使用K近邻分类，在10折交叉时的准确率，请用python为我实现。\n",
    "        \"\"\"\n",
    "        k = 1\n",
    "        knn = KNeighborsClassifier(n_neighbors=k, metric='cosine')\n",
    "\n",
    "        features = np.array(self.features)\n",
    "\n",
    "        labels = np.array(self.labels)\n",
    "\n",
    "        cv_accuracy = cross_val_score(knn, features, labels, cv=5)\n",
    "\n",
    "        for fold, accuracy in enumerate(cv_accuracy,1):\n",
    "            print(f\"Fold {fold}: {accuracy}\")\n",
    "\n",
    "        # 打印平均准确率\n",
    "        mean_accuracy = np.mean(cv_accuracy)\n",
    "        print(f\"Average Accuracy: {mean_accuracy}\")\n",
    "\n",
    "    def gather_feature_label(self,roles, roles_list):\n",
    "        features = []\n",
    "        labels = []\n",
    "        dim = 0\n",
    "\n",
    "        count = 0\n",
    "\n",
    "        for role in roles:\n",
    "            print(role,end=' ')\n",
    "\n",
    "            for file in roles_list[roles.index(role)]:\n",
    "                # print(file)\n",
    "\n",
    "                deal_flag = False\n",
    "\n",
    "                for candidate in self.candidate_path:\n",
    "\n",
    "                    candidate_fname = os.path.join(candidate,'voice',file)\n",
    "\n",
    "                    if os.path.exists(candidate_fname):\n",
    "                        # print(candidate_fname,'found')\n",
    "                        deal_flag = True\n",
    "                        feature_fname = os.path.join(candidate,'feature',file) + '.pkl'\n",
    "                        break\n",
    "\n",
    "                if deal_flag == False:\n",
    "                    print('warning!',file,'not found')\n",
    "                    continue\n",
    "\n",
    "                if not os.path.exists(feature_fname):\n",
    "                    print('warning!',feature_fname,'not found')\n",
    "                    continue\n",
    "\n",
    "                # pinkle load feature_fname\n",
    "                with open(feature_fname,'rb') as f:\n",
    "                    feature = pickle.load(f)\n",
    "\n",
    "                count += 1\n",
    "\n",
    "                # append numpy array feature into numpy matrix features\n",
    "                if dim == 0:\n",
    "                    features = feature\n",
    "                    dim = feature.shape[0]\n",
    "                    # print(dim)\n",
    "                else:\n",
    "                    features = np.vstack((features,feature))\n",
    "\n",
    "                labels.append(role)\n",
    "\n",
    "        return features, labels\n",
    "\n",
    "\n",
    "    def get_feat_sel(self):\n",
    "        roles_sel = []\n",
    "        roles_list_sel = []\n",
    "\n",
    "        M = 8\n",
    "\n",
    "        for role in self.roles[:]:\n",
    "            wav_list = self.roles_list[self.roles.index(role)]\n",
    "\n",
    "            # if len(wav_list) < M:\n",
    "                # continue\n",
    "\n",
    "            # random pick 5 element from wav_list\n",
    "            random.shuffle(wav_list)\n",
    "            # wav_list = wav_list[:]\n",
    "\n",
    "            roles_sel.append(role)\n",
    "            roles_list_sel.append(wav_list)\n",
    "\n",
    "        # print(roles)\n",
    "        # print(roles_sel)\n",
    "\n",
    "\n",
    "        feat_sel, label_sel = self.gather_feature_label(roles_sel,roles_list_sel)\n",
    "        return feat_sel, label_sel\n",
    "\n",
    "    def get_sel_predict(self):\n",
    "\n",
    "\n",
    "        corrent_dists = []\n",
    "        wrong_dists = []\n",
    "\n",
    "        for i in range(len(self.labels)):\n",
    "            # read i-th row from features, save as feat\n",
    "            feat = self.features[i, :]\n",
    "            # read i-th row from labels, save as label\n",
    "            label = self.labels[i]\n",
    "\n",
    "            # predict label of i-th row\n",
    "            predicted_label, distance = self.my_classifier.predict(feat)\n",
    "\n",
    "            # if distance < 1e-3:\n",
    "            #     continue\n",
    "\n",
    "            if label == predicted_label:\n",
    "                corrent_dists.append(distance)\n",
    "            else:\n",
    "                wrong_dists.append(distance)\n",
    "\n",
    "\n",
    "\n",
    "    def get_pridict(self):\n",
    "        threshold_certain = 0.4\n",
    "        threshold_doubt = 0.6\n",
    "        for idx,feature_folder in enumerate(self.candidate_path):\n",
    "            name = feature_folder.split('/')[-1]\n",
    "            save_name = os.path.join(self.srt_out_dir,f'{name}.txt')\n",
    "            feature_folder = os.path.join(feature_folder,\"feature\")\n",
    "\n",
    "            file_list = os.listdir(feature_folder)\n",
    "\n",
    "            N_files = len( os.listdir(feature_folder) )\n",
    "            N_files += 100\n",
    "\n",
    "            with open(save_name, \"w\", encoding=\"utf-8\") as f_out:\n",
    "\n",
    "                for id in range(N_files):\n",
    "\n",
    "                    deal_file_name = '';\n",
    "\n",
    "                    for file in file_list:\n",
    "                        if file.startswith(str(id) + '_') and file.endswith('.wav.pkl'):\n",
    "                            deal_file_name = file\n",
    "                            id_str = file.split('_')[1]\n",
    "                            break\n",
    "\n",
    "                    if deal_file_name == '':\n",
    "                        # print('file not found')\n",
    "                        continue\n",
    "\n",
    "                    full_file_name = os.path.join(feature_folder, deal_file_name)\n",
    "\n",
    "                    with open(full_file_name,'rb') as f:\n",
    "                        feature = pickle.load(f)\n",
    "\n",
    "                    predicted_label, distance = self.my_classifier.predict(feature)\n",
    "\n",
    "                    role_name = ''\n",
    "\n",
    "                    if distance < threshold_certain:\n",
    "                        role_name = predicted_label\n",
    "                    elif distance < threshold_doubt:\n",
    "                        role_name = '(可能)' + predicted_label\n",
    "\n",
    "                    output_str = role_name + ':「' +  id_str[:-8] + '」'\n",
    "\n",
    "                    # print(output_str)\n",
    "                    f_out.write(output_str + \"\\n\")\n",
    "\n",
    "                # break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2U7z1UIT-IFI"
   },
   "source": [
    "#### 运行3 台本识别"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rU_KoIxT-Ger",
    "outputId": "650a7cae-083e-4f5f-a767-cc0637e961ff"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "谷口忘了叫啥的角色朝仓春日朝比奈森小姐阿虚多丸圭一多丸裕阿虚妹妹长门古泉新川先生谷口 warning! 73_算了  不管它.wav not found\n",
      "warning! 65_我跟凉宫她初中同班三年  所以清楚得很.wav not found\n",
      "warning! 74_总之有人用它画了个很大很大  而且很怪的象形文字似的东西.wav not found\n",
      "warning! 68_没错  初中的时候她就是个令人无法理解的怪人了.wav not found\n",
      "warning! 80_是她自己亲口说的  所以错不了.wav not found\n",
      "忘了叫啥的角色 warning! 76_啊  我记得看过这个.wav not found\n",
      "朝仓 warning! 228_你可要趁着这个势头  一口气让凉宫同学跟班上的其他同学打成一片哦.wav not found\n",
      "warning! 220_我不论怎么找凉宫同学说话  她都不会理我.wav not found\n",
      "warning! 423_我会退出书法社  加入你们的社团.wav not found\n",
      "warning! 424_但是  我并不清楚文学社是在做什么的….wav not found\n",
      "warning! 420_啊  是这样啊….wav not found\n",
      "warning! 417_那个  我参加了书法社….wav not found\n",
      "warning! 229_难得有缘同班  大家不是应该好好相处吗.wav not found\n",
      "春日 warning! 377_我来介绍  她叫朝比奈实玖瑠哦.wav not found\n",
      "warning! 29_我来自……中学  请多关照.wav not found\n",
      "warning! 32_你们之中要是有外星人  未来人  异世界人或者超能力者的话.wav not found\n",
      "warning! 201_能够住在如此漂亮的别墅里  我感到非常荣幸.wav not found\n",
      "warning! 58_那就别跟我搭话  会浪费我的时间的.wav not found\n",
      "warning! 30_我毕业于东中  叫做凉宫春日.wav not found\n",
      "warning! 284_没有的话  自己组建一个就好了！.wav not found\n",
      "warning! 205_啊  是吗？.wav not found\n",
      "warning! 121_你睡什么觉啊  笨蛋.wav not found\n",
      "warning! 193_又不是外星人  未来人或者超能力者.wav not found\n",
      "warning! 206_初次见面  别墅主人～.wav not found\n",
      "朝比奈 warning! 167_啊  啊～好厉害呢～.wav not found\n",
      "森小姐 warning! 175_不  没有.wav not found\n",
      "阿虚 warning! 16_是十分渴望那些外星人  未来人  幽灵  妖怪  超能力者.wav not found\n",
      "warning! 53_呃  就是外星人什么的那些.wav not found\n",
      "warning! 229_那要是看起来就古怪的话  你又要怎么解释？.wav not found\n",
      "warning! 9_机灵的我  也早就怀疑只在圣诞节才工作的老头子是否真的存在了.wav not found\n",
      "warning! 13_不  说不定我早就发现了.wav not found\n",
      "warning! 20_不知何时开始  我就不再热心地观看电视上的UFO特别节目或灵异节目了.wav not found\n",
      "warning! 25_初中毕业之后  我也从这种小孩子的幻想中毕业了.wav not found\n",
      "warning! 10_不过  那些外星人  未来人  幽灵  妖怪  超能力者.wav not found\n",
      "warning! 18_不过  现实却是意外地残酷.wav not found\n",
      "多丸圭一 warning! 223_虽然没有自动锁  所以不会被关在门外.wav not found\n",
      "warning! 217_这是西式建筑  所以可以穿着鞋进来.wav not found\n",
      "warning! 216_别站着说话了  请进吧.wav not found\n",
      "warning! 199_哎呀  欢迎.wav not found\n",
      "warning! 305_哎呀哎呀  大家玩得很高兴呢.wav not found\n",
      "多丸裕 阿虚妹妹 warning! 76_你听我说  阿虚他刚才还打算把我赶走呢.wav not found\n",
      "长门 古泉 warning! 349_新川  有什么事吗？.wav not found\n",
      "warning! 63_换言之  名侦探所到之处必有事件发生.wav not found\n",
      "warning! 58_的确  现实世界是不会像故事那样的.wav not found\n",
      "warning! 69_也就是说呢  她想尝试一次看看.wav not found\n",
      "新川先生 warning! 360_直截了当地说  就是这么一回事.wav not found\n",
      "warning! 354_也去他可能会去的地方找过  但并没有找到.wav not found\n",
      "warning! 310_我还有工作要完成  先失陪了.wav not found\n",
      "warning! 158_的确非常不方便  请各位见谅.wav not found\n",
      "warning! 184_并没有呢  至少现在还没有.wav not found\n",
      "warning! 281_因此中午之前无法奉陪  请各位谅解.wav not found\n",
      "warning! 350_是的  刚刚我让森去裕少爷那里.wav not found\n"
     ]
    }
   ],
   "source": [
    "audio_classification = AudioClassification(audio_config['audio_roles_dir'],\n",
    "                                               srt_config['srt_out_dir'],\n",
    "                                               audio_config['audio_out_dir'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "32GPvMPN-jhi"
   },
   "outputs": [],
   "source": [
    "audio_classification.get_pridict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AJgyCWWy_EmO"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
